import os
import csv
import torch
import argparse
from tqdm import tqdm

import warnings
warnings.filterwarnings("ignore", category=Warning)

def process_nodes(data, output_path, chunk_size=50000):
    """处理大规模节点数据"""

    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    num_x_nodes = data['x'].size(0)
    num_y_nodes = data['y'].size(0)
    num_nodes = min(num_x_nodes, num_y_nodes)

    if not isinstance(data['x'], torch.Tensor):
        data['x'] = torch.tensor(data['x'])

    if data['x'].layout == torch.sparse_csr:
        data['x'] = data['x'].to_dense().contiguous()

    elif data['x'].is_sparse:
        data['x'] = data['x'].to_dense()

    with open(output_path, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(["id:ID", "label:INT"])
        
        for start in tqdm(range(0, num_nodes, chunk_size), desc="Processing Nodes"):
 
            end = min(start + chunk_size, num_nodes)
            chunk_x = data['x'][start:end].to(torch.float16)
            chunk_y = data['y'][start:end]

            if len(chunk_x) != len(chunk_y):
                break

            for i in range(len(chunk_x)):
                if i >= len(chunk_y):
                    break
                
                label = chunk_y[i].argmax().item() if chunk_y.dim() > 1 else chunk_y[i].item()
                writer.writerow([start + i, label])  

            del chunk_x, chunk_y
            torch.cuda.empty_cache() if torch.cuda.is_available() else None

if __name__ == "__main__":

    graph_name = "Yelp"

    parser = argparse.ArgumentParser(description='Process node data')

    parser.add_argument('--input', type=str, default=f'../data/{graph_name}_data.pt', 
                       help='Input data path')
    
    parser.add_argument('--output', type=str, default=f'./{graph_name}_nodes.csv',
                       help='Output CSV path')
    
    parser.add_argument('--chunk_size', type=int, default=100000,
                       help='Number of nodes per processing chunk')
    
    args = parser.parse_args()

    try:
        raw_data = torch.load(args.input)
        data = raw_data[0] if isinstance(raw_data, (list, tuple)) else raw_data
    except Exception as e:
        print(f"数据加载失败: {str(e)}")
        exit(1)
    
    process_nodes(data, args.output, args.chunk_size)

    print(f"============数据集{graph_name}的nodes.csv文件生成完毕！！！============")